package com.tuniu.agents.impl.trip;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.agent.AgentOptions;
import com.tuniu.agents.agent.OptionsAgent;
import com.tuniu.agents.common.handler.IStreamResponseHandler;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.impl.trip.enums.PlannerEnum;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.util.ChatResponseUtils;
import com.tuniu.agents.util.I18nUtil;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.beans.factory.annotation.Autowired;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;


/**
 * Planner receives four complete itineraries with resource lists and outputs recommendations
 * to clients with summaries and comparisons
 */

public class Planner extends OptionsAgent {

    public static final String AGENT_ID = "Planner";
    public static final String ITINERARY_PLAN_AGENT_ID = "ItineraryPlanner";
    private static final Logger log = LoggerFactory.getLogger(Planner.class);

    private AgentManager agentManager;
    private MemoryManager memoryManager;
    private IStreamResponseHandler streamResponseHandler;

    @Autowired
    private MessageSender messageSender;
    public Planner(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, ChatOptions chatOptions, IStreamResponseHandler streamResponseHandler) {
        super(agentManagerGroup.getModelBuilder(),
                agentManagerGroup.getChatMemory(),
                agentManagerGroup.getAdvisors().stream().map(advisor -> advisor instanceof AgentChatMemoryAdvisor ? new AgentChatMemoryAdvisor(agentManagerGroup.getChatMemory(), false, false) : advisor).collect(Collectors.toList()),
                agentManagerGroup.getPromptManager(),
                agentManagerGroup.getChatEnvironmentMemory(),
                new AgentOptions(agentId, promptTemplateId,  null, null, chatOptions));

        this.agentManager = agentManagerGroup.getAgentManager();
        this.memoryManager = agentManagerGroup.getMemoryManager();
        this.streamResponseHandler = streamResponseHandler;
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> context) {
        String message = messages.get(0).getText();
        log.info("PlannerAgent processing conversationID: {}", conversationId);
        String plannerMessage = I18nUtil.getMessage(conversationId, PlannerEnum.PLANNER.getMessageKey());
        messageSender.send(conversationId, BaseAgentNewResp.sendAnalysisMsg(plannerMessage, 1, null));


        // Get thinking content, need to check if ChatResponseUtils.getThinkTextMessage(chatResponse) exists

        AtomicBoolean thinkingChunk = new AtomicBoolean(true);

        StringBuilder thinkingContentBuffer = new StringBuilder();

        Flux<String> stringFlux = this.streamResponse(conversationId, List.of(new UserMessage(message)))
                .flatMap(response -> {
                    String thinkTextMessage = ChatResponseUtils.getThinkTextMessage(response);
                    if (StringUtils.isNotBlank(thinkTextMessage)) {
                            if (thinkingChunk.get()) {
                                String plannerAnalyseMessage = I18nUtil.getMessage(conversationId, PlannerEnum.PLANNER_ANALYSE.getMessageKey());
                                messageSender.send(conversationId, BaseAgentNewResp.sendAnalysisMsg(plannerAnalyseMessage, 1, null));
                                thinkingChunk.set(false);
                            }
                            thinkingContentBuffer.append(thinkTextMessage);

                            if(StringUtils.isNotBlank(thinkTextMessage)){
                                messageSender.send(conversationId, BaseAgentNewResp.sendThinkingMsg(null, thinkTextMessage));
                            }

                            return Flux.empty();
                        } else {
                            return Flux.just(ChatResponseUtils.getFirstTextMessage(response));
                        }
                    }
                );
        Map<String, Object> params = Map.of("message", message + "\n\n"  + thinkingContentBuffer,"conversationId", conversationId);
        return streamResponseHandler == null ? stringFlux : streamResponseHandler.handle(params, stringFlux);
    }
}
